import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

# 读取CSV文件
data = pd.read_csv('data_resource\\radiation_data.csv')
# 转换时间列为日期时间格式
data['timestamp'] = pd.to_datetime(data['timestamp'])

# 找到最早日期作为基准日期
baseline_date = data['timestamp'].min()

# 创建DataFrame
df = pd.DataFrame({'Time': data['timestamp'], 'Concentration': data['concentration']})
# 将时间戳转换为相对于基准日期的天数
df['Days_From_Baseline'] = (df['Time'] - baseline_date).dt.days

# 划分训练集和测试集
train_df, test_df = train_test_split(df, test_size=0.1, shuffle=False)
# 训练线性回归模型
model = LinearRegression()
model.fit(train_df[['Days_From_Baseline']], train_df['Concentration'])
# 预测
test_df['Predicted'] = model.predict(test_df[['Days_From_Baseline']])
# 绘图
plt.figure(figsize=(10, 6))
plt.plot(train_df['Time'], train_df['Concentration'], label='Train Data')
plt.plot(test_df['Time'], test_df['Concentration'], label='Test Data')
plt.plot(test_df['Time'], test_df['Predicted'], label='Predicted Data')
plt.xlabel('Time')
plt.ylabel('Concentration')
plt.legend()
plt.show()
